const std = @import("std");
const print = std.debug.print;
const yazap = @import("yazap");

pub const libbpf = @cImport({
    @cInclude("libbpf.h");
    @cInclude("bpf.h");
    @cInclude("btf.h");
});

var should_exit = std.atomic.Value(bool).init(false);

fn signal_handler(_: c_int) align(1) callconv(.C) void {
    std.debug.print("Received SIGINT\n", .{});

    should_exit.store(true, .seq_cst);
}

pub fn main() !void {
    var gpa = std.heap.GeneralPurposeAllocator(.{}){};
    const allocator = gpa.allocator();

    defer {
        const deinit_status = gpa.deinit();

        if (deinit_status == .leak) @panic("TEST FAIL");
    }

    var app = yazap.App.init(allocator, "pysonar4zig", "A tool to trace https stream of python");
    defer app.deinit();

    var cmd = app.rootCommand();

    try cmd.addArg(yazap.Arg.singleValueOption("libssl_path", 'l', "libssl path"));

    const matches = try app.parseProcess();

    if (matches.getSingleValue("libssl_path") == null) {
        print("Please specify libssl path.\n\n\n", .{});
        try app.displayHelp();
        return error.NO_LIBSSL_PATH;
    }
    const libssl_path = matches.getSingleValue("libssl_path").?;

    var sa = std.posix.Sigaction{
        .handler = .{
            .handler = &signal_handler,
        },
        .mask = std.posix.empty_sigset,
        .flags = 0,
    };
    std.posix.sigaction(std.posix.SIG.INT, &sa, null);
    std.posix.sigaction(std.posix.SIG.TERM, &sa, null);

    const bytes = @embedFile("@bpf_prog");

    const obj = libbpf.bpf_object__open_mem(bytes.ptr, bytes.len, null);
    if (obj == null) {
        print("failed to open bpf object: {}\n", .{std.posix.errno(-1)});
        return error.OPEN;
    }
    defer libbpf.bpf_object__close(obj);

    const ret = libbpf.bpf_object__load(obj);
    if (ret != 0) {
        print("failed to load bpf object: {}\n", .{std.posix.errno(ret)});
        return error.LOAD;
    }

    const prog = libbpf.bpf_object__find_program_by_name(obj, "uprobe_SSL_write_ex").?;

    const opts = bpf_uprobe_opts{ .func_name = "SSL_write_ex" };

    const link = libbpf.bpf_program__attach_uprobe_opts(prog, -1, @ptrCast(libssl_path), 0, @ptrCast(&opts)) orelse {
        print("failed to attach entry_prog {s}: {}\n", .{ libbpf.bpf_program__name(prog), std.posix.errno(-1) });
        return error.ATTACH;
    };
    defer _ = libbpf.bpf_link__destroy(link);

    const events = libbpf.bpf_object__find_map_by_name(obj, "events").?;

    const ring_buf = libbpf.ring_buffer__new(libbpf.bpf_map__fd(events), on_sample, null, null).?;
    defer libbpf.ring_buffer__free(ring_buf);

    while (!should_exit.load(.seq_cst)) {
        const ret2 = libbpf.ring_buffer__poll(ring_buf, 100);
        if (ret2 == -4) {
            break;
        }
        if (ret2 < 0) {
            print("failed to poll ring buffer: {}\n", .{std.posix.errno(ret)});
            break;
        }
    }

    print("Exiting gracefully...\n", .{});
}

const Event = extern struct {
    pid: u32,
    uid: u32,
    comm: [40]u8,
    buf: [3000]u8,
};

fn on_sample(_: ?*anyopaque, _data: ?*anyopaque, _: usize) callconv(.C) c_int {
    const c: *Event = @ptrCast(@alignCast(_data.?));

    const pid = c.pid;
    const uid = c.uid;
    const comm = c.comm;
    const buf = c.buf;

    print("CMD [{s}] PID [{d}], UID [{d}] =========>\n{s}\n<=========\n", .{ comm, pid, uid, buf });

    return 0;
}

const bpf_uprobe_opts = extern struct {
    sz: u64 = @sizeOf(bpf_uprobe_opts),
    ref_ctr_offset: u64 = 0,
    bpf_cookie: u64 = 0,
    retprobe: bool = false,
    func_name: [*]const u8,
    attach_mode: u32 = 0,
};
